import torch
import numpy
import torch.nn as nn
import argparse
import os
from PIL import Image
from torch.utils.data import DataLoader
from dataset import TestDataset
from models import DFN
os.environ['CUDA_VISIBLE_DEVICES']='1'


def save_results(preds, num_class, result_dir, image_name):
    COLOR_MAP = [(128, 64, 128), (244, 35, 232), (70, 70, 70), (102, 102, 156), (190, 153, 153), (153, 153, 153),
                 (250, 170, 30), (220, 220, 0), (107, 142, 35), (152, 251, 152), (70, 130, 180), (220, 20, 60),
                 (255, 0, 0), (0, 0, 142), (0, 0, 70), (0, 60, 100), (0, 80, 100), (0, 0, 230), (119, 11, 32)]
    for i in range(len(preds)):
        pre = preds[i]
        row, col = pre.shape
        dst = torch.zeros((row, col, 3))
        for j in range(num_class):
            dst[pre == j] = torch.Tensor(COLOR_MAP[j])
        dst = numpy.array(dst, dtype=numpy.uint8)
        dst = Image.fromarray(dst)
        dst = dst.resize((dst.size[0]*4, dst.size[1]*4), Image.NEAREST)
        dst.save(result_dir+'/'+image_name[i], 'PNG')

def main():
    batch_size = 1
    start_cuda = 'cuda:0'
    gpu_ids = [0]

    num_class = 19
    ckpt_path = './ckpt'
    result = './result'
    save_result = True

    encoder = 'shufflenet'#encoder = 'resnet' #
    net_weight_path = 'shufflenet-dfn-v1/model_epoch_5.pth' #net_weight_path = 'resnet-dfn-v1/model_epoch_5.pth' #

    dataset=TestDataset('test',  '.png')
    testloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    input()
    model_id = encoder+'-dfn-v1'
    result_dir = result+'/'+model_id
    if not os.path.isdir(result_dir):
        os.makedirs(result_dir)

    net = DFN(num_class, weights=None, encoder=encoder)


    device = torch.device(start_cuda if torch.cuda.is_available() else "cpu")
    if torch.cuda.device_count() > 1:
        net = nn.DataParallel(net, device_ids=gpu_ids)
    if net_weight_path:
        print('loading the net weight : {}'.format(net_weight_path))
        net.load_state_dict(torch.load(ckpt_path + '/' + net_weight_path))
        net.eval()
    net.to(device)

    # In[ ]:

    #training
    print('start testing on GPU : {}'.format(gpu_ids))

    for i, data in enumerate(testloader, 0):
        inputs, _, image_name = data
        inputs = inputs.float()
        #         labels = labels.float()
        inputs = inputs.to(device)
        outputs = net(inputs)
        preds = torch.argmax(outputs[-2].data, 1)
        if save_result:
            save_results(preds, num_class, result_dir, image_name)


if __name__ == '__main__':
    main()